#define _CRT_SECURE_NO_WARNINGS
#define HAVE_REMOTE
#define WIN32

#include <stdio.h>
#include <stdlib.h>
#include <pcap.h>
#include <Packet32.h>
#include <ntddndis.h>
#include <string.h>
#include <time.h>
#include <winsock2.h> // For WSAStartup, WSACleanup, inet_ntoa, ntohs, htons
#include <ctype.h>    // For isprint

#pragma comment(lib, "Packet")
#pragma comment(lib, "wpcap")
#pragma comment(lib, "WS2_32")

// --- Constants ---
#define LOG_FILENAME "ftp_login_log.csv"
#define FTP_CONTROL_PORT 21
#define MAX_PAYLOAD_CHECK 256 // Limit how much payload we inspect
#define MAX_LOGIN_INFO_LEN 128 // Max length for username/password storage
#define MAX_TRACKED_CONNECTIONS 1024 // Max concurrent connections to track (simple approach)

// --- Data Structures ---

// Ethernet Header (14 bytes) - same as before
typedef struct mac_header {
    u_char dest_addr[6];
    u_char src_addr[6];
    u_short type; // Network byte order
} mac_header;

// IP Header - same as before (using struct in_addr)
typedef struct ip_header {
    u_char ver_ihl;
    u_char tos;
    u_short tlen;
    u_short identification;
    u_short flags_fo;
    u_char ttl;
    u_char proto;
    u_short crc;
    struct in_addr saddr;
    struct in_addr daddr;
} ip_header;

// TCP Header (minimum 20 bytes)
typedef struct tcp_header {
    u_short sport; // Source port (network byte order)
    u_short dport; // Destination port (network byte order)
    u_int seq;     // Sequence number (network byte order)
    u_int ack_seq; // Acknowledgment number (network byte order)
    u_char data_offset; // Data offset (4 bits) + Reserved (4 bits)
    u_char flags;       // Control flags (CWR, ECE, URG, ACK, PSH, RST, SYN, FIN)
    u_short window;     // Window size (network byte order)
    u_short checksum;   // Checksum (network byte order)
    u_short urg_ptr;    // Urgent pointer (network byte order)
    // Options follow here...
} tcp_header;

// Simple structure to temporarily store login info per connection (using source IP/Port as key)
typedef struct {
    struct in_addr src_ip;
    u_short src_port; // Network byte order
    char last_user[MAX_LOGIN_INFO_LEN];
    char last_pass[MAX_LOGIN_INFO_LEN];
    time_t last_seen; // Timestamp to clear old entries
} FtpLoginAttempt;

// --- Global Variables ---
FILE* log_file_ftp = NULL;
FtpLoginAttempt login_attempts[MAX_TRACKED_CONNECTIONS]; // Simple fixed-size array for tracking
int login_attempt_count = 0;
time_t cleanup_threshold = 60; // Cleanup entries older than 60 seconds

// --- Function Prototypes ---
void ftp_packet_handler(u_char* param, const struct pcap_pkthdr* header, const u_char* pkt_data);
void find_or_create_attempt(struct in_addr sip, u_short sport, FtpLoginAttempt** attempt);
void record_ftp_log(const char* ts, const char* smac, const char* sip, const char* dmac, const char* dip,
                    const char* user, const char* pass, const char* result);
void cleanup_old_attempts();
void format_mac(const u_char* mac, char* buffer, size_t buffer_size); // Re-use from previous
void write_ftp_log_header();
int is_first_ftp_log = 1;

// --- Main Function ---
int main() {
    // ... (WSAStartup, Device Discovery, Selection - same as previous code) ...
    int inum;
    pcap_if_t* alldevs;
    pcap_if_t* d;
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t* adhandle;
    int i = 0;
    u_int netmask;
    // **** Filter changed to capture TCP traffic on port 21 ****
    char packet_filter[] = "tcp and port 21";
    struct bpf_program fcode;
    WSADATA wsaData;

    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) { /* ... error handling ... */ return -1; }
    if (pcap_findalldevs_ex(/* ... */) == -1) { /* ... error handling ... */ WSACleanup(); exit(1); }
    for (d = alldevs; d; d = d->next) { /* ... list devices ... */ }
    if (i == 0) { /* ... no devices found ... */ WSACleanup(); return -1; }
    printf("Enter the interface number (1-%d):", i);
    if (scanf("%d", &inum) != 1 || inum < 1 || inum > i) { /* ... invalid input ... */ WSACleanup(); return -1; }
    while (getchar() != '\n'); // Clear buffer
    for (d = alldevs, i = 0; i < inum - 1; d = d->next, i++);
    if ((adhandle = pcap_open(/* ... */)) == NULL) { /* ... error handling ... */ WSACleanup(); return -1; }
    if (pcap_datalink(adhandle) != DLT_EN10MB) { /* ... error handling ... */ WSACleanup(); return -1; }
    if (d->addresses != NULL && d->addresses->netmask != NULL) netmask = ((struct sockaddr_in*)(d->addresses->netmask))->sin_addr.s_addr;
    else netmask = PCAP_NETMASK_UNKNOWN;
    if (pcap_compile(adhandle, &fcode, packet_filter, 1, netmask) < 0) { /* ... error handling ... */ WSACleanup(); return -1; }
    if (pcap_setfilter(adhandle, &fcode) < 0) { /* ... error handling ... */ pcap_freecode(&fcode); WSACleanup(); return -1; }
    pcap_freecode(&fcode);


    // --- Initialize FTP Log File ---
    log_file_ftp = fopen(LOG_FILENAME, "a");
    if (log_file_ftp == NULL) {
        perror("Error opening FTP log file");
        pcap_close(adhandle);
        pcap_freealldevs(alldevs);
        WSACleanup();
        return -1;
    }
     fseek(log_file_ftp, 0, SEEK_END);
     if (ftell(log_file_ftp) == 0) {
         write_ftp_log_header();
     } else {
         is_first_ftp_log = 0;
     }
     fflush(log_file_ftp);

    // Initialize the simple tracking array
    memset(login_attempts, 0, sizeof(login_attempts));
    login_attempt_count = 0;

    printf("\nListening for FTP login attempts on %s (Port 21)...\nPress Ctrl+C to stop.\n", d->description);
    pcap_freealldevs(alldevs);

    // --- Start Capture Loop (using ftp_packet_handler) ---
    pcap_loop(adhandle, 0, ftp_packet_handler, NULL);

    // --- Cleanup ---
    printf("\nCapture finished.\n");
    if (log_file_ftp != NULL) {
        fclose(log_file_ftp);
    }
    pcap_close(adhandle);
    WSACleanup();
    return 0;
}

// --- FTP Packet Handler Callback ---
void ftp_packet_handler(u_char* param, const struct pcap_pkthdr* header, const u_char* pkt_data) {
    mac_header* mh;
    ip_header* ih;
    tcp_header* tcph;
    const u_char* payload;
    unsigned int ip_header_len, tcp_header_len, payload_len;
    struct tm* ltime;
    char timestr[TIMESTAMP_STR_LEN];
    char src_mac_str[MAC_ADDR_STR_LEN], dst_mac_str[MAC_ADDR_STR_LEN];
    char src_ip_str[IP_ADDR_STR_LEN], dst_ip_str[IP_ADDR_STR_LEN];
    FtpLoginAttempt* current_attempt = NULL;
    int is_client_to_server = 0; // Flag to check direction based on port 21

    // --- Timestamp ---
    time_t local_tv_sec = header->ts.tv_sec;
    ltime = localtime(&local_tv_sec);
    if (ltime != NULL) strftime(timestr, sizeof(timestr), "%Y-%m-%d %H:%M:%S", ltime);
    else strncpy(timestr, "Time Error", sizeof(timestr)-1), timestr[sizeof(timestr)-1]='\0';


    // --- Parse Headers ---
    mh = (mac_header*)pkt_data;
    if (ntohs(mh->type) != 0x0800) return; // Not an IP packet

    ih = (ip_header*)(pkt_data + sizeof(mac_header));
    ip_header_len = (ih->ver_ihl & 0x0F) * 4;
    if (ip_header_len < 20 || header->caplen < (sizeof(mac_header) + ip_header_len)) return; // Malformed/Truncated IP
    if (ih->proto != IPPROTO_TCP) return; // Not a TCP packet

    tcph = (tcp_header*)((u_char*)ih + ip_header_len);
    tcp_header_len = ((tcph->data_offset >> 4) & 0x0F) * 4; // TCP header length in bytes
    if (tcp_header_len < 20 || header->caplen < (sizeof(mac_header) + ip_header_len + tcp_header_len)) return; // Malformed/Truncated TCP

    // --- Calculate Payload ---
    payload = (u_char*)tcph + tcp_header_len;
    payload_len = ntohs(ih->tlen) - ip_header_len - tcp_header_len;
    if ((long)payload_len < 0) return; // Invalid length calculation
    // Ensure payload length doesn't exceed captured length
    if (sizeof(mac_header) + ip_header_len + tcp_header_len + payload_len > header->caplen) {
        payload_len = header->caplen - (sizeof(mac_header) + ip_header_len + tcp_header_len);
        if ((long)payload_len < 0) payload_len = 0;
    }


    // --- Extract Basic Info ---
    format_mac(mh->src_addr, src_mac_str, sizeof(src_mac_str));
    format_mac(mh->dest_addr, dst_mac_str, sizeof(dst_mac_str));
    strncpy(src_ip_str, inet_ntoa(ih->saddr), sizeof(src_ip_str) - 1); src_ip_str[sizeof(src_ip_str)-1] = '\0';
    strncpy(dst_ip_str, inet_ntoa(ih->daddr), sizeof(dst_ip_str) - 1); dst_ip_str[sizeof(dst_ip_str)-1] = '\0';

    // --- Identify Direction and Find/Create Login Attempt State ---
    if (ntohs(tcph->dport) == FTP_CONTROL_PORT) { // Client sending to server
        is_client_to_server = 1;
        find_or_create_attempt(ih->saddr, tcph->sport, ¤t_attempt);
    } else if (ntohs(tcph->sport) == FTP_CONTROL_PORT) { // Server sending to client
        is_client_to_server = 0;
        // Use DESTINATION IP/Port from server's perspective (which is client's IP/Port)
        find_or_create_attempt(ih->daddr, tcph->dport, ¤t_attempt);
    } else {
        return; // Not related to FTP control port 21
    }

    if (!current_attempt) {
        // fprintf(stderr, "Warning: Could not track connection state (max connections reached?).\n");
        return; // Max connections reached or allocation failed
    }
     current_attempt->last_seen = local_tv_sec; // Update last seen time

    // --- Parse FTP Commands/Responses (Simplified) ---
    if (payload_len > 0 && payload != NULL) {
        // Limit inspection length for safety and simplicity
        unsigned int check_len = payload_len > MAX_PAYLOAD_CHECK ? MAX_PAYLOAD_CHECK : payload_len;
        char payload_str[MAX_PAYLOAD_CHECK + 1];

        // Copy payload safely and null-terminate
        // Only copy printable characters, replace others with '.'
        for (unsigned int k=0; k < check_len; ++k) {
            payload_str[k] = isprint(payload[k]) ? payload[k] : '.';
        }
        payload_str[check_len] = '\0';


        // --- Client Commands ---
        if (is_client_to_server) {
            if (strncmp(payload_str, "USER ", 5) == 0) {
                char* user_start = payload_str + 5;
                char* user_end = strstr(user_start, "\r\n"); // Find CRLF
                if (user_end) {
                    size_t user_len = user_end - user_start;
                    if (user_len < MAX_LOGIN_INFO_LEN) {
                        strncpy(current_attempt->last_user, user_start, user_len);
                        current_attempt->last_user[user_len] = '\0';
                        current_attempt->last_pass[0] = '\0'; // Clear previous password on new USER
                        // printf("Debug: Found USER '%s' for %s:%d\n", current_attempt->last_user, src_ip_str, ntohs(tcph->sport));
                    }
                }
            } else if (strncmp(payload_str, "PASS ", 5) == 0) {
                 if (current_attempt->last_user[0] != '\0') { // Only record PASS if USER was seen
                    char* pass_start = payload_str + 5;
                    char* pass_end = strstr(pass_start, "\r\n");
                    if (pass_end) {
                        size_t pass_len = pass_end - pass_start;
                        if (pass_len < MAX_LOGIN_INFO_LEN) {
                            strncpy(current_attempt->last_pass, pass_start, pass_len);
                            current_attempt->last_pass[pass_len] = '\0';
                            // printf("Debug: Found PASS '****' for user '%s' from %s:%d\n", current_attempt->last_user, src_ip_str, ntohs(tcph->sport));
                        }
                    }
                 }
            }
        }
        // --- Server Responses ---
        else { // Server to client
            if (strncmp(payload_str, "230 ", 4) == 0) { // Login successful
                if (current_attempt->last_user[0] != '\0' && current_attempt->last_pass[0] != '\0') {
                    record_ftp_log(timestr, src_mac_str, src_ip_str, dst_mac_str, dst_ip_str,
                                   current_attempt->last_user, current_attempt->last_pass, "SUCCEED");
                    // Clear state after successful/failed log
                    current_attempt->last_user[0] = '\0';
                    current_attempt->last_pass[0] = '\0';
                }
            } else if (strncmp(payload_str, "530 ", 4) == 0) { // Login incorrect
                if (current_attempt->last_user[0] != '\0' && current_attempt->last_pass[0] != '\0') {
                     record_ftp_log(timestr, src_mac_str, src_ip_str, dst_mac_str, dst_ip_str,
                                   current_attempt->last_user, current_attempt->last_pass, "FAILED");
                     // Clear state after successful/failed log
                     current_attempt->last_user[0] = '\0';
                     current_attempt->last_pass[0] = '\0';
                } else if (current_attempt->last_user[0] != '\0') {
                     // Handle case where 530 might occur after USER but before PASS (e.g., user unknown)
                     // Log with empty password? Or wait for PASS? This simple model logs failure if USER was seen.
                      record_ftp_log(timestr, src_mac_str, src_ip_str, dst_mac_str, dst_ip_str,
                                   current_attempt->last_user, "", "FAILED"); // Log fail with empty pass
                     current_attempt->last_user[0] = '\0';
                     current_attempt->last_pass[0] = '\0';
                }
            }
            // Could also check for 331 (Password required) to confirm USER was accepted
        }
    }

     // --- Cleanup old tracked connections ---
     cleanup_old_attempts();
}

// --- Helper Functions ---

// Find existing or create new FtpLoginAttempt based on client IP/Port
void find_or_create_attempt(struct in_addr sip, u_short sport, FtpLoginAttempt** attempt) {
    *attempt = NULL;
    time_t now = time(NULL);

    // Try to find existing entry
    for (int i = 0; i < login_attempt_count; ++i) {
        // Compare IP and Port (network byte order comparison for port is fine here)
        if (login_attempts[i].src_ip.s_addr == sip.s_addr && login_attempts[i].src_port == sport) {
            *attempt = &login_attempts[i];
            (*attempt)->last_seen = now; // Update timestamp
            return;
        }
    }

    // Not found, try to add a new one if space available
    if (login_attempt_count < MAX_TRACKED_CONNECTIONS) {
        int index = login_attempt_count++;
        login_attempts[index].src_ip = sip;
        login_attempts[index].src_port = sport;
        login_attempts[index].last_user[0] = '\0';
        login_attempts[index].last_pass[0] = '\0';
        login_attempts[index].last_seen = now;
        *attempt = &login_attempts[index];
        return;
    }

    // No space, try to find the oldest entry to overwrite (simple LRU approximation)
    int oldest_index = -1;
    time_t oldest_time = now;
    for (int i = 0; i < MAX_TRACKED_CONNECTIONS; ++i) {
         if (login_attempts[i].last_seen < oldest_time) {
              oldest_time = login_attempts[i].last_seen;
              oldest_index = i;
         }
    }
    if (oldest_index != -1) {
         // Overwrite the oldest entry
         login_attempts[oldest_index].src_ip = sip;
         login_attempts[oldest_index].src_port = sport;
         login_attempts[oldest_index].last_user[0] = '\0';
         login_attempts[oldest_index].last_pass[0] = '\0';
         login_attempts[oldest_index].last_seen = now;
         *attempt = &login_attempts[oldest_index];
         // fprintf(stderr,"Warning: Overwriting old connection state for %s:%d\n", inet_ntoa(sip), ntohs(sport));
    } else {
         // Should not happen if MAX_TRACKED_CONNECTIONS > 0
         *attempt = NULL;
    }

}

// Record the FTP login event to the CSV file
void record_ftp_log(const char* ts, const char* smac, const char* sip, const char* dmac, const char* dip,
                    const char* user, const char* pass, const char* result) {
    if (log_file_ftp != NULL) {
        // Basic sanitization: replace commas in user/pass with semicolons to avoid breaking CSV
        char sanitized_user[MAX_LOGIN_INFO_LEN];
        char sanitized_pass[MAX_LOGIN_INFO_LEN];
        strncpy(sanitized_user, user, sizeof(sanitized_user) - 1); sanitized_user[sizeof(sanitized_user)-1]='\0';
        strncpy(sanitized_pass, pass, sizeof(sanitized_pass) - 1); sanitized_pass[sizeof(sanitized_pass)-1]='\0';
        for(char *p = sanitized_user; *p; ++p) if(*p == ',') *p = ';';
        for(char *p = sanitized_pass; *p; ++p) if(*p == ',') *p = ';';

        fprintf(log_file_ftp, "%s,%s,%s,%s,%s,%s,%s,%s\n",
                ts, smac, sip, dmac, dip,
                sanitized_user, sanitized_pass, result);
        fflush(log_file_ftp); // Ensure write
    }
}

// Simple cleanup: remove entries not seen recently (not very robust)
void cleanup_old_attempts() {
     time_t now = time(NULL);
     int write_index = 0;
     for (int read_index = 0; read_index < login_attempt_count; ++read_index) {
          if (now - login_attempts[read_index].last_seen <= cleanup_threshold) {
               if (write_index != read_index) {
                    login_attempts[write_index] = login_attempts[read_index];
               }
               write_index++;
          }
     }
     login_attempt_count = write_index; // Update the count of active entries
}

// Format MAC address (same as before)
void format_mac(const u_char* mac, char* buffer, size_t buffer_size) {
    snprintf(buffer, buffer_size, "%02X-%02X-%02X-%02X-%02X-%02X",
             mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
}

// Write the CSV header for FTP log
void write_ftp_log_header() {
    if (log_file_ftp != NULL && is_first_ftp_log) {
        fprintf(log_file_ftp, "Timestamp,Source MAC,Source IP,Destination MAC,Destination IP,Username,Password,Status\n");
        fflush(log_file_ftp);
        is_first_ftp_log = 0;
    }
}